using System;
using System.Collections;
using System.Text;

using Mapack;

namespace A4G.Core
{
	public class Viterbi
	{
		private struct StateMapping
		{
			public StateMapping(double probability, int[] v_path, double v_prob)
			{
				this.probability = probability;
				this.v_path = v_path;
				this.v_prob = v_prob;
			}

			public double probability;
			public int[] v_path;
			public double v_prob;
		}

		public char[] CalculateMostProbableStatePath(char[] sequence, char[] observations, char[] states, Matrix a, Matrix e, double[] initialStateProbabilities)
		{
			if (states.Length != initialStateProbabilities.Length)
			{
				throw new ApplicationException("Inconsistent number of states and initial state probabilities.");
			}

			if ((states.Length != a.Rows) || (states.Length != a.Columns))
			{
				throw new ApplicationException("Inconsistent number of states and state transition probabilities.");
			}

			if (states.Length != e.Rows)
			{
				throw new ApplicationException("Inconsistent number of states and emission probabilities.");
			}

			if (observations.Length != e.Columns)
			{
				throw new ApplicationException("Inconsistent number of observations and emission probabilities.");
			}

			StateMapping[] T = new StateMapping[states.Length];
			for (int stateIndex = 0; stateIndex < states.Length; stateIndex++)
			{
				T[stateIndex] = new StateMapping(
					initialStateProbabilities[stateIndex],
					new int[] { stateIndex },
					initialStateProbabilities[stateIndex]);
			}

			double total = 0;
			ArrayList maxpath = new ArrayList();
			double valmax = double.MinValue;

			for (int observationIndex = 0; observationIndex < sequence.Length; observationIndex++)
			{
				int output = GetIndexOf(observations, sequence[observationIndex]);

				StateMapping[] U = new StateMapping[states.Length];
				for (int nextStateIndex = 0; nextStateIndex < states.Length; nextStateIndex++)
				{
					total = 0;
					valmax = double.MinValue;
					for (int currentStateIndex = 0; currentStateIndex < states.Length; currentStateIndex++)
					{
						double p = e[currentStateIndex, output] * a[currentStateIndex, nextStateIndex];

						StateMapping current = T[currentStateIndex];
						current.probability *= p;
						current.v_prob *= p;
						total += current.probability;

						if (current.v_prob > valmax)
						{
							maxpath.Clear();
							maxpath.AddRange(current.v_path);
							maxpath.Add(nextStateIndex);
							valmax = current.v_prob;
						}
					}
					U[nextStateIndex] = new StateMapping(total, (int[]) maxpath.ToArray(typeof(int)), valmax);
				}
				T = U;
			}

			total = 0;
			maxpath.Clear();
			valmax = double.MinValue;
			for (int currentStateIndex = 0; currentStateIndex < states.Length; currentStateIndex++)
			{
				StateMapping current = T[currentStateIndex];
				total += current.probability;

				if (current.v_prob > valmax)
				{
					maxpath.Clear();
					maxpath.AddRange(current.v_path);
					maxpath.Add(currentStateIndex);
					valmax = current.v_prob;
				}
			}

			ArrayList mostProbablyStatePath = new ArrayList();
			foreach (int stateIndex in (int[])maxpath.ToArray(typeof(int)))
			{
				mostProbablyStatePath.Add(states[stateIndex]);
			}

			return (char[]) mostProbablyStatePath.ToArray(typeof(char));
		}

		private int GetIndexOf(char[] chars, char token)
		{
			for (int i = 0; i < chars.Length; i++)
			{
				if (chars[i].Equals(token))
				{
					return i;
				}
			}

			throw new ApplicationException("Token could not be found in character array.");
		}
	}
}
